package main

import (
	"crypto/tls"
	"encoding/csv"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

var tr = &http.Transport{
	TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
}
var httpClient = &http.Client{Transport: tr}

func main() {
	argsNum := len(os.Args)
	var csvFile = "D:/document/畅遇/obs-migration/img2.txt"
	if argsNum > 1 {
		csvFile = os.Args[1]
	}
	if csvFile == "" {
		fmt.Println("指定路径")
		return
	}
	if _, err := os.Stat(csvFile); os.IsNotExist(err) {
		fmt.Println("文件不存在")
		return
	}
	outputDir := filepath.Dir(csvFile) + "/images"

	maxConcurrentDownloads := 10

	err := downloadImages(csvFile, outputDir, maxConcurrentDownloads)
	if err != nil {
		fmt.Println("Error:", err)
	}
}

func downloadImages(csvFile, outputDir string, maxConcurrentDownloads int) error {
	// Create the output directory if it doesn't exist
	if _, err := os.Stat(outputDir); os.IsNotExist(err) {
		err = os.Mkdir(outputDir, 0755)
		if err != nil {
			return err
		}
	}

	// Open the CSV file
	file, err := os.Open(csvFile)
	if err != nil {
		return err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	var wg sync.WaitGroup
	semaphore := make(chan struct{}, maxConcurrentDownloads)

	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}

		url := strings.Trim(record[0], "'")

		wg.Add(1)
		semaphore <- struct{}{} // acquire a slot in the semaphore
		go func(id, url string) {
			defer func() {
				wg.Done()
				<-semaphore // release the slot in the semaphore
			}()

			// Download the image
			err := downloadImage(id, url, outputDir)
			if err != nil {
				fmt.Printf("Error downloading image for ID %s: %v\n", id, err)
			}
		}("", url)
	}

	wg.Wait()
	return nil
}

func downloadImage(id, url, outputDir string) error {
	response, err := httpClient.Get(url)
	if err != nil {
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP request failed with status code: %d", response.StatusCode)
	}

	// Get the file extension from the URL
	//fileExtension := filepath.Ext(url)
	basename := filepath.Base(url)

	// Create the output file
	outputFile := filepath.Join(outputDir, basename)
	file, err := os.Create(outputFile)
	if err != nil {
		return err
	}
	defer file.Close()

	// Copy the image data to the output file
	_, err = io.Copy(file, response.Body)
	if err != nil {
		return err
	}

	fmt.Printf("Downloaded image for basename %s\n", basename)
	return nil
}
